/*
  Copyright (c) 2010 Davide Rossetti (davide.rossetti@roma1.infn.it),
  Andrea Maiorano (andrea.maiorano@roma1.infn.it)

  This file is part of CuHB (CUDA Heisenberg) package.
  
  CuHB is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  CuHB is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  
  You should have received a copy of the GNU General Public License
  along with CuHB.  If not, see <http://www.gnu.org/licenses/>.

*/

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>

#include <cuda_runtime.h>

#include "hb.h"
#include "rng.hpp"
#include "algos.hpp"

//-----------------------------------------------------------------------------

static float cpu_calc_energy_r1(const Svec& v, const Jmat&j)
{
    assert(1 == v.m_replicas);

    const int lx = v.m_cfg.lat[DIR_X];
    const int ly = v.m_cfg.lat[DIR_Y];
    const int lz = v.m_cfg.lat[DIR_Z];

    float en = 0.0f;

    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                vec lf = make_vec(.0f, .0f, .0f);
                lf += j(ix,iy,iz,DIR_X) * v(mod(ix+1,lx),iy,iz);
                lf += j(ix,iy,iz,DIR_Y) * v(ix,mod(iy+1,ly),iz);
                lf += j(ix,iy,iz,DIR_Z) * v(ix,iy,mod(iz+1,lz));
                //int ixm = mod(ix-1,lx); lf += j(ixm,iy,iz,DIR_X) * v(ixm,iy,iz);
                //int iym = mod(iy-1,ly); lf += j(ix,iym,iz,DIR_Y) * v(ix,iym,iz);
                //int izm = mod(iz-1,lz); lf += j(ix,iy,izm,DIR_Z) * v(ix,iy,izm);
                float den = lf * v(ix,iy,iz);
                en += den;
#if 0
                printf("den[%d,%d,%d]=%f\n", ix, iy, iz, den);
                jmat jj = j(ix,iy,iz,DIR_X);
                printf("j={ %f %f %f | %f %f %f | %f %f %f }\n",
                       jj.j00, jj.j01, jj.j02,
                       jj.j10, jj.j11, jj.j12,
                       jj.j20, jj.j21, jj.j22);
                vec vv = v(ix,iy,iz);
                printf("v=(%f,%f,%f)\n", vv.x, vv.y, vv.z);
#endif
            }

    return -en/(lx*ly*lz);
}

//-----------------------------------------------------------------------------

static inline size_t block(size_t x, size_t y, size_t z)
{
    int i = XYZ2IDX(x,y,z);
    return i/BLOCK_SIZE;
}

//-----------------------------------------------------------------------------

static void cpu_heatbath_r1_v1(float beta, rng_uniform &rng, Svec& res, const Svec& v, const Jmat&j)
{
    assert(1 == v.m_replicas);

    const int lx = v.m_cfg.lat[DIR_X];
    const int ly = v.m_cfg.lat[DIR_Y];
    const int lz = v.m_cfg.lat[DIR_Z];

    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                float2 r = rng.gen_uniform_float2(block(ix,iy,iz));
                vec lf = {.0f, .0f, .0f};
                lf += j(ix,iy,iz,DIR_X) * v(mod(ix+1,lx),iy,iz);
                lf += j(ix,iy,iz,DIR_Y) * v(ix,mod(iy+1,ly),iz);
                lf += j(ix,iy,iz,DIR_Z) * v(ix,iy,mod(iz+1,lz));
                int ixm = mod(ix-1,lx); lf += j(ixm,iy,iz,DIR_X) * v(ixm,iy,iz);
                int iym = mod(iy-1,ly); lf += j(ix,iym,iz,DIR_Y) * v(ix,iym,iz);
                int izm = mod(iz-1,lz); lf += j(ix,iy,izm,DIR_Z) * v(ix,iy,izm);
                vec h = beta * lf;
                float h2 = h * h;
                float mh = std::sqrt(h2);
                float inv_mh = 1/mh;
#if 1
                double expf = exp(2*mh);
                // float expf = std::exp(2*beta*mh); NaN here!!
                float logf = log(1 + r.x*(expf-1));
                float costheta = logf*inv_mh-1;
                float ctf = (std::abs(1-costheta)>1.0e-6) ? 0 : 1; // 1 if costheta~1
                float sintheta = std::sqrt(1 - costheta*costheta + ctf) * (1 - ctf);
#else
                float costheta = std::log(1+r.x*(std::exp(2*mh)-1))/mh-1;
                float sintheta = std::sqrt(1 - costheta*costheta);
#endif

                vec hnorm = inv_mh * h;
                float mhxy = std::sqrt(1 - hnorm.z*hnorm.z);
                float inv_mhxy = 1.0f/mhxy;
                vec c = make_vec(hnorm.x*inv_mhxy, hnorm.y*inv_mhxy, hnorm.z);
                float cosphi, sinphi;
                float phi = 2 * M_PI *r.y;
                sincos(phi, &sinphi, &cosphi);
                vec rr = make_vec(c.x * (cosphi * sintheta * c.z + costheta * mhxy) - sintheta * sinphi * c.y,
                                  c.x * sintheta * sinphi + c.y * (cosphi * c.z * sintheta + costheta * mhxy),
                                  costheta * c.z - cosphi * sintheta * mhxy);
                if(is_nan(rr.x) || is_nan(rr.y) || is_nan(rr.z))
                    info("nan!! [%d,%d,%d]={%f,%f,%f} h2=%f r1=%f r2=%f mh=%f costheta=%f sintheta=%f\n", 
                         ix, iy, iz, 
                         rr.x, rr.y, rr.z, 
                         h2, r.x, r.y, mh, costheta, sintheta);
                res.set(ix,iy,iz, rr);
            }
}

//-----------------------------------------------------------------------------

static void cpu_heatbath_r1_v2(float beta, rng_uniform &rng, Svec& res, const Svec& v, const Jmat&j)
{
    assert(1 == v.m_replicas);

    const int lx = v.m_cfg.lat[DIR_X];
    const int ly = v.m_cfg.lat[DIR_Y];
    const int lz = v.m_cfg.lat[DIR_Z];

    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                float2 r = rng.gen_uniform_float2();
                vec lf = {.0f, .0f, .0f};
                lf += j(ix,iy,iz,DIR_X) * v(mod(ix+1,lx),iy,iz);
                lf += j(ix,iy,iz,DIR_Y) * v(ix,mod(iy+1,ly),iz);
                lf += j(ix,iy,iz,DIR_Z) * v(ix,iy,mod(iz+1,lz));
                int ixm = mod(ix-1,lx); lf += j(ixm,iy,iz,DIR_X) * v(ixm,iy,iz);
                int iym = mod(iy-1,ly); lf += j(ix,iym,iz,DIR_Y) * v(ix,iym,iz);
                int izm = mod(iz-1,lz); lf += j(ix,iy,izm,DIR_Z) * v(ix,iy,izm);
                vec h = lf;
                float h2 = h * h;
                float mh = std::sqrt(h2);
                float inv_v = 1.0/std::sqrt(lf.x*lf.x + lf.y*lf.y);
                float inv_mh = 1/mh;
                double expf = std::exp((double)2*beta*mh);
                // float expf = std::exp(2*beta*mh); NaN here!!
                float logf = std::log(1 + r.x*(expf-1));
                float costheta = logf/beta*inv_mh-1;
                float ctf = (std::abs(1-costheta)>1.0e-6) ? 0 : 1; // 1 if costheta~1
                float sintheta = std::sqrt(1 - costheta*costheta + ctf) * (1 - ctf);
                if(is_nan(sintheta))
                    info("[%d,%d,%d] costheta=%f(0x%08x) sintheta=%f(0x%08x)\n", 
                         ix, iy, iz, costheta, cast_as_hex(costheta), sintheta, cast_as_hex(sintheta));

                vec hnorm = inv_mh * h;
                float2 v = make_float2(lf.y*inv_v, -lf.x*inv_v);
                vec ns = make_vec(hnorm.x*costheta - v.y*hnorm.z*sintheta,
                                  hnorm.y*costheta + v.x*hnorm.z*sintheta,
                                  hnorm.z*costheta + (hnorm.z*v.y-hnorm.y*v.x)*sintheta);
                float cosphi, sinphi;
                sincos(2 * M_PI *r.y, &sinphi, &cosphi);
                if(r.y > 0.5)
                    sinphi *= -1.0;
                float fac = costheta*(1-cosphi);
                // la=x mu=y nu=z
                vec rr = make_vec(ns.x * cosphi + fac * hnorm.x + (ns.y*hnorm.z-ns.z*hnorm.y)*sinphi,
                                  ns.y * cosphi + fac * hnorm.y + (ns.z*hnorm.x-ns.x*hnorm.z)*sinphi,
                                  ns.z * cosphi + fac * hnorm.z + (ns.x*hnorm.y-ns.y*hnorm.x)*sinphi);
                float rnorm2 = rr * rr;
                if(is_nan(rr.x) || is_nan(rr.y) || is_nan(rr.z) || (std::abs(1-rnorm2)>1.0e-5))
                    info("nan!! [%d,%d,%d] "
                         "lfield={%f,%f,%f} "
                         "ns={%f,%f,%f} "
                         "h2=%f r1=%f r2=%f mh=%f "
                         "costheta=%f sintheta=%f "
                         "res={%f,%f,%f} res2=%f\n", 
                         ix, iy, iz, 
                         lf.x, lf.y, lf.z,
                         ns.x, ns.y, ns.z,
                         h2, r.x, r.y, mh, 
                         costheta, sintheta,
                         rr.x, rr.y, rr.z, rnorm2);
                res.set(ix,iy,iz, rr);
            }
}

//-----------------------------------------------------------------------------

#define cpu_heatbath_r1 cpu_heatbath_r1_v1

static const vec init_dir = make_vec(1.0, 0.0, 0.0);

static void check_spins(Svec &s)
{
    //s.dump();

    float a = s.calc_alignment(init_dir);
    if(std::abs(a-1)> 1.0e-5)
        warn("alignment=%f check failed\n", a);

    if(!s.check_norm())
        warn("1-norm spin check failed\n");
}

static void do_basic_test(Config &cfg, rng_uniform &rng)
{
    size_t pad = 0;
    Svec spins(cfg, pad);
    Svec res(cfg, pad);
    Jmat j(cfg, pad);

    size_t nspins = cfg.lat[DIR_X]*cfg.lat[DIR_Y]*cfg.lat[DIR_Z];

    res.init(0.0, 0.0, 0.0);

    printf("# init spins to (%f,%f,%f)\n", init_dir.x, init_dir.y, init_dir.z);
    spins.init(init_dir);
    spins.check_norm();
    //spins.cpu_to_gpu();
    check_spins(spins);

    printf("# init J to ident=diagonal(1,1,1)\n");
    j.init_ident();
    //printf("init J to Gaussian\n");
    //j.init_gaussian(rng);
    //j.cpu_to_gpu();

    float energy = 0.0;
    energy = cpu_calc_energy_r1(spins, j);
    printf("# energy=%f\n", energy);

    float temp = 0.1f;
    float beta = 1/temp;

    size_t nsteps = NITERS;
    size_t nprints = nsteps/10;
    printf("# executing %d HB steps at beta=%f\n", nsteps, beta);
    for(int k=0; k<nsteps; k+=2) {
        cpu_heatbath_r1(beta, rng, res, spins, j);
#if 0
        printf("dumping 'res'\n");
        res.dump();
        energy = cpu_calc_energy_r1(res, j);
        printf("energy 'res'=%f\n", energy);        
        //check_spins(res);
#endif
        cpu_heatbath_r1(beta, rng, spins, res, j);
#if 0
        printf("dumping 'spins'\n");
        spins.dump();
        energy = cpu_calc_energy_r1(spins, j);
        printf("energy 'res'=%f\n", energy);        
        //check_spins(spins);
#endif
#if 1
        if((k+2) % nprints == 0) {
            energy = cpu_calc_energy_r1(spins, j);
            check_spins(spins);
            printf("%d %f\n", k+2, energy);
        }
#endif
    }
    
    energy = cpu_calc_energy_r1(spins, j);
    printf("# final energy %f\n", energy);
}

//-----------------------------------------------------------------------------

void benchmark_cpu_hb3d_r1()
{
    Config cfg;
    cfg.dims = D;
    cfg.lat[DIR_X] = LX;
    cfg.lat[DIR_Y] = LY;
    cfg.lat[DIR_Z] = LZ;
    cfg.cpu_precision = SinglePrecision;
    cfg.gpu_precision = SinglePrecision;

    const int niters = NITERS;

    size_t pad = 0;
    Svec spins(cfg, pad);
    Svec res(cfg, pad);
    Jmat j(cfg, pad);
    float e;
    float beta = 5.0f;

    size_t nspins = cfg.lat[DIR_X]*cfg.lat[DIR_Y]*cfg.lat[DIR_Z];
    size_t n_blocks = (nspins*spins.replicas())/BLOCK_SIZE;
    printf("# %s\n", __FUNCTION__);
    printf("# Heisenberg %dD:\n", cfg.dims);
    printf("#  lattice:          %dx%dx%d\n", cfg.lat[DIR_X], cfg.lat[DIR_Y], cfg.lat[DIR_Z]);
    printf("#  replicas:         %u\n", spins.replicas());
    printf("#  tot spins:        %d\n", nspins);
    printf("#  precision:        %s\n", "single");
    printf("#  num iterations:   %d\n", niters);
    printf("#  thr block size:   %d\n", BLOCK_SIZE);
    printf("#  num thr blocks:   %d\n", n_blocks);
    printf("#  spin vector size: %.3fKB\n", spins.m_gpu_bytes/(1000.0));
    printf("#  J matrix size:    %.3fKB\n", j.m_gpu_bytes/(1000.0));

    rng_uniform rng(BLOCK_SIZE, n_blocks);
    rng.init(0xf4543243, false);

    do_basic_test(cfg, rng);
}

//-----------------------------------------------------------------------------
//-----------------------------------------------------------------------------

/*
 * Local variables:
 *  mode: c++
 *  c-indent-level: 4
 *  c-basic-offset: 4
 *  tab-width: 4
 *  indent-tabs-mode: nil
 * End:
 */
